
use tokio::io::{AsyncRead, AsyncWrite};
use super::{WebSocketFrame, WebSocketReceiveError, WebSocketSendError};

pub enum WebSocketRole {
    Client,
    Server,
}

pub struct WebSocketReadStream<R> {
    reader: R,
}

pub struct WebSocketWriteStream<W> {
    writer: W,
    role: WebSocketRole,
}

impl<R: AsyncRead + Unpin> WebSocketReadStream<R> {
    pub fn new(reader: R) -> Self {
        Self { reader }
    }

    pub async fn read_frame(&mut self) -> Result<Option<WebSocketFrame>, WebSocketReceiveError> {
        let frame =  WebSocketFrame::read_from(&mut self.reader).await?;
        Ok(frame)
    }

    pub async fn receive_text(&mut self) -> Result<Option<String>, WebSocketReceiveError> {
        let frame =  WebSocketFrame::read_from(&mut self.reader).await?;
        match frame {
            Some(frame) => Ok( Some(frame.into_string()?) ),
            None => Ok(None),
        }
    }

    pub async fn receive_binary(&mut self) -> Result<Option<Vec<u8>>, WebSocketReceiveError> {
        let frame =  WebSocketFrame::read_from(&mut self.reader).await?;
        match frame {
            Some(frame) => Ok( Some(frame.into_bytes()?) ),
            None => Ok(None),
        }
    }
}

impl<W: AsyncWrite + Unpin> WebSocketWriteStream<W> {
    pub fn new(writer: W, role: WebSocketRole) -> Self {
        Self { writer, role }
    }

    pub async fn send_frame(&mut self, frame: &WebSocketFrame) -> Result<(), WebSocketSendError> {
        frame.send_to(&mut self.writer).await
    }

    pub async fn send_text<S: Into<String>>(&mut self, text: S) -> Result<(), WebSocketSendError> {
        let frame = match self.role {
            WebSocketRole::Client => WebSocketFrame::client_text(text),
            WebSocketRole::Server => WebSocketFrame::server_text(text),
        };
        self.send_frame(&frame).await
    }

    pub async fn send_binary<B: Into<Vec<u8>>>(&mut self, bytes: B) -> Result<(), WebSocketSendError> {
        let frame = match self.role {
            WebSocketRole::Client => WebSocketFrame::client_binary(bytes),
            WebSocketRole::Server => WebSocketFrame::server_binary(bytes),
        };
        self.send_frame(&frame).await
    }

    pub async fn send_close(&mut self) -> Result<(), WebSocketSendError> {
        let frame = match self.role {
            WebSocketRole::Client => WebSocketFrame::client_close(),
            WebSocketRole::Server => WebSocketFrame::server_close(),
        };
        self.send_frame(&frame).await
    }
}